#include "Basis.h"
#include "Common.h"

BOOLEAN ProcessIoCtlForCommon(ULONG ioControlCode, PDEVICE_OBJECT deviceObject, PIRP irp, PIO_STACK_LOCATION irpStack)
{
	UNREFERENCED_PARAMETER(deviceObject);
	UNREFERENCED_PARAMETER(irpStack);

	switch (ioControlCode)
	{
	case IOCTL_GET_KPCR_ADR:
		ProcessGetKpcrAdr(irp);
		break;
	case IOCTL_READ_MEMORY:
		ProcessReadMemory(irp);
		break;
	case IOCTL_WRITE_MEMORY:
		ProcessWriteMemory(irp);
		break;
	case IOCTL_INVOKE_FUNCTION:
		ProcessInvokeFunction(irp);
		break;
	default:
		return FALSE;
	}
	return TRUE;
}

VOID ProcessGetKpcrAdr(PIRP irp)
{
	// get kpcr address
	ULONG_PTR pKpcr;
	KeSetSystemAffinityThread(1);
	__asm
	{
		mov eax, fs:[0x1C]; // SelfPcr
		mov pKpcr, eax;
	}
	KeRevertToUserAffinityThread();

	if (pKpcr == 0)
	{
		irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
		return;
	}

	// save to buffer
	PULONG_PTR buffer = (PULONG_PTR)irp->AssociatedIrp.SystemBuffer;
	*buffer = pKpcr;
	irp->IoStatus.Information = sizeof(ULONG_PTR);
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID ProcessReadMemory(PIRP irp)
{
	PREAD_MEMORY_REQUEST pReadMemoryReuqest = (PREAD_MEMORY_REQUEST)irp->AssociatedIrp.SystemBuffer;
	if (pReadMemoryReuqest->Address == 0 || pReadMemoryReuqest->Length == 0 || !MmIsAddressValid((PVOID)pReadMemoryReuqest->Address))
	{
		irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
		return;
	}

	ULONG Length = pReadMemoryReuqest->Length; // will overwrite so copy its value
	KdPrint(("reading address 0x%x, length %d", pReadMemoryReuqest->Address, Length));
	RtlCopyMemory(irp->AssociatedIrp.SystemBuffer, (PVOID)pReadMemoryReuqest->Address, Length);
	irp->IoStatus.Information = Length;
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID ProcessWriteMemory(PIRP irp)
{
	PWRITE_MEMORY_REQUEST pWriteMemoryRequest = (PWRITE_MEMORY_REQUEST)irp->AssociatedIrp.SystemBuffer;
	if (pWriteMemoryRequest->Address == 0 || pWriteMemoryRequest->Length == 0 || !MmIsAddressValid((PVOID)pWriteMemoryRequest->Address))
	{
		irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
		return;
	}

	KdPrint(("writing address 0x%x, length %d", pWriteMemoryRequest->Address, pWriteMemoryRequest->Length));
	RtlCopyMemory((PVOID)pWriteMemoryRequest->Address, &pWriteMemoryRequest->Value, pWriteMemoryRequest->Length);
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID ProcessInvokeFunction(PIRP irp)
{
	PINVOKE_FUNCTION_REQUEST pInvokeFunctionRequest = (PINVOKE_FUNCTION_REQUEST)irp->AssociatedIrp.SystemBuffer;
	PVOID pFunc = GetFunctionAddess(pInvokeFunctionRequest->FunctionName);
	if (pFunc == NULL)
	{
		irp->IoStatus.Status = STATUS_NOT_FOUND;
		return;
	}

	__try
	{
		// pushing parameters
		for (LONG parameterCount = pInvokeFunctionRequest->ParamCount - 1; parameterCount >= 0; parameterCount--)
		{
			ULONG_PTR Parameter = pInvokeFunctionRequest->Params[parameterCount];
			_asm push Parameter
		}

		// call function
		ULONG_PTR returnValue;
		_asm
		{
			call pFunc
			mov returnValue, eax
		}
		if (pInvokeFunctionRequest->ReturnSize > 0)
		{
			if (pInvokeFunctionRequest->IsReturnValueType)
			{
				*(PULONG_PTR)irp->AssociatedIrp.SystemBuffer = returnValue;
				irp->IoStatus.Information = sizeof(ULONG_PTR);
			}
			else
			{
				ULONG returnSize = pInvokeFunctionRequest->ReturnSize; // will overwrite so copy its value
				RtlCopyMemory(irp->AssociatedIrp.SystemBuffer, (PVOID)returnValue, returnSize);
				irp->IoStatus.Information = returnSize;
			}
		}
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
	}
}

PVOID GetFunctionAddess(PCWSTR FunctionName)
{
	UNICODE_STRING UniCodeFunctionName;
	RtlInitUnicodeString(&UniCodeFunctionName, FunctionName);
	return MmGetSystemRoutineAddress(&UniCodeFunctionName);
}
